# coding=utf-8
import numpy as np  # type: ignore
import matplotlib  # type: ignore

matplotlib.use('TkAgg')
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt  # type: ignore
import rclpy
from rclpy.node import Node
from rclpy.timer import Timer
import rclpy.qos
from scipy.spatial.transform import Rotation as R  # type: ignore
from std_msgs.msg import Float64MultiArray
import sys, select, os
import time
import math
import socket
import struct
from geometry_msgs.msg import Twist, PoseStamped, Pose
from interfaces_consensus.msg import CommandMsg
from builtin_interfaces.msg import Time
from geometry_msgs.msg import Pose
from geometry_msgs.msg import Point, Quaternion

EARTH_RADIUS = 6371008.7714150598


class traj_server(Node):
    def __init__(self):
        # 初始化节点，指定节点名称
        super().__init__('traj_server_virtual')
        self.folnum = 1  # 无人机数量

        self.original_position = [0, 0, 0]

        self.uav_id = int(sys.argv[1])  #无人机id
        self.cluster_list = [] #集群组成列表
        self.childCluster_list = [] #子集群组成列表


        self.flag = True  #用于发送第三个点的判断，避免无人机停顿
        self.task_flag = True  #用于散开执行任务判断   
        # self.finish_replan_flag = True #用于对齐执行列表和规划结果
        self.command_flag = True

        self.longitude = 0.0
        self.latitude = 0.0
        self.altitude = 0.0
        self.uavpos_x = 0.0
        self.uavpos_y = 0.0
        self.uavpos_z = 0.0

        self.security = True

        self.timestamps = [None]

        # self.uav_pose_pubs = {}

        self.receive_traj = False
        self.receive_ori = False

        self.colony_list = list(range(1, self.folnum + 1))  # 生成 1 到 folnum 的列表
        # print(f"colony_list={self.colony_list}")
        self.cache_colony_list = list(range(1, self.folnum + 1)) 

        # 动态生成 Pose 对象列表，每个 Pose 对象存储无人机的当前位置
        self.poses = [Float64MultiArray() for _ in range(self.folnum)]
        topic_plan_cmd_pub = "/drone{}/planning/command".format(self.uav_id)
        self.plan_cmd_sub = self.create_subscription(
            Float64MultiArray,  # 消息类型
            topic_plan_cmd_pub,  # 话题名称
            self.planning_command_callback,  # 回调函数
            10
        )
        topic_ori_pose_pub = "/drone{}/original_position".format(self.uav_id)
        self.ori_pose_sub = self.create_subscription(
            Float64MultiArray,  # 消息类型
            topic_ori_pose_pub,  # 话题名称
            self.original_position_callback,  # 回调函数
            10
        )
        topic_name = "/drone{}/velocity_and_position".format(self.uav_id)
        self.get_logger().info(f"Listen to {topic_name}")

        #凌云测试使用
        self.publisher_ = self.create_publisher(Float64MultiArray, topic_name, 10)

        self.create_subscription(
            Float64MultiArray,
            topic_name,
            lambda msg, id=self.uav_id: self.pose_callback(msg, id),
            10
        )

        qos_profile_2 = rclpy.qos.QoSProfile(
            reliability=rclpy.qos.ReliabilityPolicy.RELIABLE,
            history=rclpy.qos.HistoryPolicy.KEEP_LAST,
            depth=10
        )

        self.timer1 = self.create_timer(0.1, self.cmdCallback)

        # for i in range(self.folnum):
        #     index = i + 1
        #     topic = "/drone{}/setpoint_raw/local".format(index)
        #     if i not in self.uav_pose_pubs:
        #         self.uav_pose_pubs[i] = self.create_publisher(Float64MultiArray, topic, 10)

        topic = "/drone{}/setpoint_raw/local".format(self.uav_id)
        # if i not in self.uav_pose_pubs:
        self.uav_pose_pubs = self.create_publisher(Float64MultiArray, topic, 10)

        cmd_topic = "/drone{}/consensus_APF/command".format(self.uav_id)
        self.cmd_pub = self.create_publisher(CommandMsg, cmd_topic, qos_profile_2)
        self.subscription = self.create_subscription(
            CommandMsg,  # 消息类型
            cmd_topic,  # 话题名称
            self.consensus_APF_command_callback,  # 回调函数
            qos_profile_2
        )

        # # UDP点对点发送
        # self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        # self.server_address = ('192.168.2.230', 5007)
        # # self.sock.bind(self.server_address) # 绑定套接字到指定地址和端口,发送不需要这句
        # UDP多播发送
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM, socket.IPPROTO_UDP)
        self.sock.setsockopt(socket.IPPROTO_IP, socket.IP_MULTICAST_TTL, 32)

    def cmdCallback(self):
        # 给无人机发送位置和朝向角
        if self.receive_traj and self.receive_ori:
            num = self.traj_pose_x.shape[0]
            if num == 2:
                self.colony_list = list([self.uav_id]) 
                self.folnum = 1
            traj_k = self.traj_pose_x.shape[1]
            if traj_k > 3 and self.flag:
                self.K = 2
                self.flag = False
            print(f"len: {traj_k}")
            for uav_i in range(self.folnum):
                longitude = self.original_position[0] + (self.traj_pose_x[uav_i, self.K] / (
                        EARTH_RADIUS + self.original_position[2])) * 180 / math.pi
                latitude = self.original_position[1] + (self.traj_pose_y[uav_i, self.K] / (
                        EARTH_RADIUS + self.original_position[2])) * 180 / math.pi
                altitude = self.original_position[2] + self.traj_pose_z[uav_i, self.K]
                velocity_east = self.traj_V_x[uav_i, self.K]
                velocity_north = self.traj_V_y[uav_i, self.K]
                velocity_up = self.traj_V_z[uav_i, self.K]

                # 判断指令位置与当前位置差是否超过3米
                x = (longitude - self.longitude) / 180 * math.pi * (
                        EARTH_RADIUS + self.altitude)
                y = (latitude - self.latitude) / 180 * math.pi * (
                        EARTH_RADIUS + self.altitude)
                z = altitude - self.altitude
                distance = np.sqrt(x ** 2 + y ** 2 + z ** 2)
                #飞大飞机时，要确定是否需要高度，现在无人机高度漂移太大
                # distance = np.sqrt(x ** 2 + y ** 2)
                if self.colony_list[uav_i] == self.uav_id:
                    # if distance < 3.0:
                    #凌云测试使用
                    if distance < 10000.0:
                        self.poses[uav_i].data = [longitude, latitude, altitude, velocity_east, velocity_north, velocity_up]
                        # print(f"executing x: {self.traj_pose_x[uav_i, self.K]}--y: {self.traj_pose_y[uav_i, self.K]}--z: {self.traj_pose_z[uav_i, self.K]}")
                        print(f"executing lon: {longitude}--lat: {latitude}--alt: {altitude}")
                        print(f"Distance to instruction x: {x}--y: {y}--z: {z}")

                        #凌云测试使用
                        # msg = Float64MultiArray()
                        # msg.data = [8.8, longitude, latitude, altitude, velocity_east, velocity_north, velocity_up]
                        # self.publisher_.publish(msg)
                        type = b'\x01'
                        head = b'\x14'
                        uav_id = self.uav_id
                        status_data = struct.pack(
                                '=2ci6d',
                                type, head, uav_id, longitude, latitude, altitude, velocity_east, velocity_north, velocity_up
                            )

                        self.sock.sendto(status_data, (f'224.0.0.1', 2025))

                        
                        
                    else:
                        
                        self.security = False
                        print("超出3m安全距离")
                if self.uav_id == self.colony_list[0]:
                    # if self.colony_list[uav_i] != self.uav_id:
                        #回传凌云系统
                        type = b'\x02'
                        head = b'\x02'
                        uav_id = self.colony_list[uav_i]
                        longitude = longitude
                        latitude = latitude
                        altitude = altitude
                        x_vel = velocity_east
                        y_vel = velocity_north
                        z_vel = velocity_up
                        #计算姿态角
                        yaw ,pitch, roll = self.calculate_angles(x_vel, y_vel, z_vel)
                        status_data = struct.pack(
                                '=2ci9d',
                                type, head, uav_id, longitude, latitude, altitude, x_vel, y_vel, z_vel, roll, pitch, yaw
                            )
                        self.sock.sendto(status_data, (f'224.1.11.111', 5007))

            # 发布对应的 poses
            for i in range(self.folnum):
                if self.colony_list[i] == self.uav_id:
                    self.uav_pose_pubs.publish(self.poses[i])

            if self.K < traj_k - 1:
                if self.security:
                    self.K += 1
                else:
                    self.security = True
            # else:
            #     if traj_k <12 and self.task_flag:
            #         self.task_flag = False
            #         msg = CommandMsg()
            #         msg.stamp = self.get_clock().now().to_msg()
            #         msg.frame_id = "task"
            #         self.cmd_pub.publish(msg)
            #         self.cache_colony_list = list([self.uav_id])
                    # self.folnum = len(self.colony_list)
                        
                        # self.sendStatus()

            print(f".............{self.K+1}")

            # plt.ion()  # 打开交互模式
            # # fig = plt.figure()  # 创建一个新的图形窗口
            # area = self.compute_area(self.traj_pose_x[7, self.K-1], self.traj_pose_y[7, self.K-1], self.traj_pose_z[7, self.K-1], 10) #暂时修改为二维
            # ArrowLength = 0.5  # 箭头长度
            # plt.ion()  # 打开交互模式
            # plt.clf()
            # for j in range(self.folnum):
            #     # plt.quiver(self.traj_pose_x[j, self.K], self.traj_pose_y[j, self.K], ArrowLength * np.cos(pose_th[j, k + 1]),
            #     #         ArrowLength * np.sin(pose_th[j, k + 1]), scale=1, scale_units='xy', color='k',
            #     #         width=0.01)  # 绘制机器人的方向
            #     if j == 9:
            #         state = 2  # 领导者标记
            #     else:
            #         state = 1  # 跟随者标记
            #     self.draw_circle(self.traj_pose_x[j, self.K -1], self.traj_pose_y[j, self.K-1], 0.25, state)  # 绘制机器人
            #     # for i in range(N):
            #     # if A[i, j] == 1:
            #     # draw_arrow([pose_x[j, k+1], pose_y[j, k+1]], [pose_x[i, k+1], pose_y[i, k+1]], 0.2)  # 绘制连接箭头，显示通信或控制关系

            #     # if len(ob_temp) != 0:
            #     #     plt.plot(ob_temp[:, 0], ob_temp[:, 1], 'Xk', linewidth=2)  # 绘制障碍物位置

            # plt.axis(area)  # 设置显示区域
            # plt.grid(True)  # 显示网格
            # plt.draw()  # 实时更新图形
            # plt.pause(0.01)  # 添加短暂的暂停以显示图形
            # # plt.show()
            # # arrive_time+=1
            # # if arrive_time%3==0 and path_times<path_len-1:
            # # path_times+=1
            # plt.ioff()  # 关闭交互模式

            # self.receive_traj = False

    def planning_command_callback(self, msg):
        # if self.finish_replan_flag:
            self.K = 0
            self.colony_list = self.cache_colony_list
            self.folnum = len(self.colony_list)
            self.poses = [Float64MultiArray() for _ in range(self.folnum)]
            # 获取行列信息
            rows = msg.layout.dim[0].size
            cols = msg.layout.dim[1].size

            # 重新组织数据为原来的矩阵
            data = np.array(msg.data).reshape(6, rows, cols)  # 6 个矩阵，每个矩阵的维度是 rows x cols

            # 提取每个矩阵
            self.traj_V_x = data[0]
            self.traj_V_y = data[1]
            self.traj_V_z = data[2]
            self.traj_pose_x = data[3]
            self.traj_pose_y = data[4]
            self.traj_pose_z = data[5]

            # 打印矩阵数据
            # self.get_logger().info(f"Received traj_V_x:\n{self.traj_V_x}")
            # self.get_logger().info(f"Received traj_V_y:\n{self.traj_V_y}")
            # self.get_logger().info(f"Received traj_V_z:\n{self.traj_V_z}")
            # self.get_logger().info(f"Received traj_pose_x:\n{self.traj_pose_x}")
            # self.get_logger().info(f"Received traj_pose_y:\n{self.traj_pose_y}")
            # self.get_logger().info(f"Received traj_pose_z:\n{self.traj_pose_z}")
            self.receive_traj = True
            self.command_flag = True
        #     self.finish_replan_flag = False
        #     if not self.command_flag and not self.finish_replan_flag:
        #         self.command_flag = True
        #         self.finish_replan_flag = True
        # else:
        #     self.get_logger().info(
        #         f'Not finish replan!!!')


    def calculate_angles(self, V_E, V_N, V_U):
        yaw = math.atan2(V_E, V_N)
        yaw = math.radians(90)-yaw

        if yaw < 0:
            yaw  += 2 * math.pi

        horizontal_speed=math.sqrt(V_E**2 + V_N**2)
        pitch = math.atan2(V_U,horizontal_speed)

        roll =0
        return yaw, pitch, roll

    def pose_callback(self, msg, drone_id):
        # self.get_logger().info(f"Drone_{drone_id+1} is getting position!")
        # 从消息中读取位置数据
        # 与数仿调试
        # if self.flag:
        #     if drone_id ==7:
        #         return
        self.longitude = msg.data[1]  # 假设单位为度*1E7
        self.latitude = msg.data[2]
        self.altitude = msg.data[3]
        self.timestamps[0] = msg.data[0]  # 记录时间戳
        self.uavpos_x = (self.longitude - self.original_position[0]) / 180 * math.pi * (
                EARTH_RADIUS + self.original_position[2])
        self.uavpos_y = (self.latitude - self.original_position[1]) / 180 * math.pi * (
                EARTH_RADIUS + self.original_position[2])
        self.uavpos_z = self.altitude - self.original_position[2]
        # print(f"get Drone_{drone_id+1} lon: {self.longitude}--lat: {self.latitude}--alt: {self.altitude}")
        # print(f"get Drone_{drone_id+1} x: {self.uavpos_x[drone_id]}--y: {self.uavpos_y[drone_id]}--alt: {self.uavpos_z[drone_id]}")

    def original_position_callback(self, msg):
        self.original_position = msg.data
        # self.get_logger().info(f"get original_position:{self.original_position}")
        self.receive_ori = True

    def sendStatus(self):
        '''集群到位信息，回传凌云，凌云控制虚拟uav散开'''
        type = b'\x03'
        head = b'\x03'
        uav_id = self.uav_id
        for i in range(3):
            num = i +1
            status_data = struct.pack(
                    '=2c2i',
                    type, head, num, uav_id)
            # print(len(status_data))
            # print(status_data)
            self.sock.sendto(status_data, (f'224.1.11.111', 5007))

    def consensus_APF_command_callback(self, msg):
        # 命令接收，包含分组信息
        if msg.frame_id =="grouping":
            print(f"grouping  {msg}")
            if self.command_flag:
                self.cache_colony_list = msg.dynamic_array
                self.childCluster_list = msg.dynamic_array
                # self.colony_list = msg.dynamic_array
                # self.folnum = len(self.colony_list)
                self.poses = [Float64MultiArray() for _ in range(self.folnum)]
                self.task_flag = True
                self.command_flag = False

                # 输出数组，查看提取的非零值
                self.get_logger().info(
                    f'Received colony_list: {self.cache_colony_list}')
            else:
                self.get_logger().info(
                    f'Not finish replan!!!')
            # self.folnum = len(self.colony_list)
        elif msg.frame_id =="formation":
            print(f"formation{msg}")
            if self.command_flag:
                self.command_flag = False

            else:
                self.get_logger().info(
                    f'Not finish replan!!!')

        elif msg.frame_id =="return_to_home":
            print(f"return_to_home{msg}")
            if self.command_flag:
                self.command_flag = False
                self.task_flag = False
                self.cache_colony_list = self.cluster_list

            else:
                self.get_logger().info(
                    f'Not finish replan!!!')
                
        elif msg.frame_id =="cluster":
            print(f"cluster{msg}")
            self.cache_colony_list = msg.dynamic_array
            self.cluster_list = msg.dynamic_array
            # self.colony_list = msg.dynamic_array
            # self.folnum = len(self.colony_list)
            self.poses = [Float64MultiArray() for _ in range(self.folnum)]

            # 输出数组，查看提取的非零值
            self.get_logger().info(
                f'Received colony_list: {self.cache_colony_list}')
            # self.folnum = len(self.colony_list)


    def draw_circle(self, x, y, r, state):
        ''' 画圆形的函数'''
        color = ''
        line = 1
        if state == 1:
            color = 'b'
        elif state == 2:
            color = 'r'
            line = 2
        elif state == 3:
            color = 'g'  # 添加额外的颜色选择
            line = 1
        elif state == 4:
            color = 'c'  # 添加额外的颜色选择
            line = 1
        elif state == 5:
            color = 'm'  # 添加额外的颜色选择
            line = 1
        elif state == 6:
            color = 'y'  # 添加额外的颜色选择
            line = 1
        elif state == 7:
            color = 'k'  # 添加额外的颜色选择
            line = 1
        t = np.arange(0, 2 * np.pi, 0.01)
        X = x + r * np.cos(t)
        Y = y + r * np.sin(t)
        plt.plot(X, Y, color=color, linewidth=line)

    def compute_area(self, x, y, z, range=10):
        '''计算动画显示的区域'''
        if x >= 0:
            num_x = np.floor(x / range)
            x1 = num_x * range
            x2 = (num_x + 1) * range
        else:
            num_x = np.floor(x / -range)
            x1 = -range * (num_x + 1)
            x2 = -range * num_x
        if y >= 0:
            num_y = np.floor(y / range)
            y1 = num_y * range
            y2 = (num_y + 1) * range
        else:
            num_y = np.floor(y / -range)
            y1 = -range * (num_y + 1)
            y2 = -range * num_y
        if z >= 0:
            num_z = np.floor(z / range)
            z1 = num_z * range
            z2 = (num_z + 1) * range
        else:
            num_z = np.floor(z / -range)
            z1 = -range * (num_z + 1)
            z2 = -range * num_z
        # area = [x1 - 8, x2 + 8, y1 - 8, y2 + 8, z1 -8, z2 + 8]
        area = [x1 - 200, x2 + 200, y1 - 200, y2 + 200]
        return area


def main(args=None):
    # 初始化 ROS 2 客户端库
    rclpy.init(args=args)

    # 创建并运行 consensus_APF_uav 实例
    node = traj_server()

    # # 使用 MultiThreadedExecutor 来支持并行回调
    # executor = MultiThreadedExecutor()
    # executor.add_node(node)

    # 让节点保持运行并处理回调函数
    # executor.spin()
    rclpy.spin(node)

    # 销毁节点并关闭 ROS 2 客户端
    # executor.shutdown()
    node.destroy_node()
    rclpy.shutdown()


if __name__ == '__main__':
    main()
